#include <cassert>

#include "common.h"
#include "cuda-play/utils/time_cost.h"
#include "cuda.h"

DEFINE_int32(N, 2048, "");
DEFINE_int32(BLOCK_SIZE, 32, "");

namespace perf {
namespace sharemem {
int N;
int BLOCK_SIZE;

using TimeCost = cuda_play::utils::TimeCost;
template <typename T>
__global__ void matrix_transpose_naive(T *input, T *output, int N) {
  int indexX = threadIdx.x + blockIdx.x * blockDim.x;
  int indexY = threadIdx.y + blockIdx.y + blockDim.y;
  int index = indexX * N + indexY;

  int transposedIndex = indexX * N + indexY;
  output[index] = input[transposedIndex];
}

template <typename T>
void cpuTranspose(T *input, T *output) {
  for (int i = 0; i < N; ++i) {
    for (int j = 0; j < N; ++j) {
      output[j * N + i] = input[i * N + j];
    }
  }
}

template <typename T>
void gpuTranspose(T *input, T *output) {
  T *dInput, *dOutput;
  size_t size = sizeof(T) * N * N;
  cudaMalloc(&dInput, size);
  cudaMalloc(&dOutput, size);
  cudaMemcpy(dInput, input, size, cudaMemcpyHostToDevice);

  dim3 blockSize(N, N, 1);
  dim3 gridSize(N / BLOCK_SIZE, N / BLOCK_SIZE, 1);
  TimeCost tc;
  matrix_transpose_naive<<<gridSize, blockSize>>>(input, output, N);
  LOG(INFO) << "pure kernel execution time: " << tc.get_elapsed();
  cudaMemcpy(output, dOutput, size, cudaMemcpyDeviceToHost);
}

TEST_F(ShareMemTest, matrix_transpose_test) {
  N = FLAGS_N;
  BLOCK_SIZE = FLAGS_BLOCK_SIZE;

  int *a = new int[N * N], *b = new int[N * N], *c = new int[N * N],
      *d = new int[N * N];

  LOG(INFO) << "N = " << N << ", BLOCK_SIZE = " << BLOCK_SIZE;

  TimeCost tc;
  cpuTranspose(a, b);
  LOG(INFO) << "cpu time used: " << tc.get_elapsed();
  gpuTranspose(c, d);
  LOG(INFO) << "gpu time used: " << tc.get_elapsed();

  for (int i = 0; i < N * N; ++i) {
    if (b[i] != d[i]) {
      LOG(INFO) << "b[i] != d[i]! "
                << "i = " << i << ", b[i] = " << b[i] << ", d[i] = " << d[i];
    }
    EXPECT_EQ(b[i], d[i]);
  }

  delete[] a;
  delete[] b;
  delete[] c;
  delete[] d;
}

}  // namespace sharemem
}  // namespace perf
